// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <pollux/common/io/io_statistics.h>

#include <pollux/common/file/file.h>

#include <gtest/gtest.h>

namespace kumo::pollux::dwio::common {

// Testing stream producing deterministic data. The byte at offset is
// the low byte of 'seed_' + offset.
class TestReadFile : public pollux::ReadFile {
 public:
  TestReadFile(
      uint64_t seed,
      uint64_t length,
      std::shared_ptr<filesystems::File::IoStats> ioStats)
      : seed_(seed), length_(length), ioStats_(std::move(ioStats)) {}

  uint64_t size() const override {
    return length_;
  }

  std::string_view pread(
      uint64_t offset,
      uint64_t length,
      void* buffer,
      filesystems::File::IoStats* stats = nullptr) const override {
    const uint64_t content = offset + seed_;
    const uint64_t available = std::min(length_ - offset, length);
    int fill;
    for (fill = 0; fill < available; ++fill) {
      reinterpret_cast<char*>(buffer)[fill] = content + fill;
    }
    if (stats) {
      stats->addCounter(
          "read", RuntimeCounter(fill, RuntimeCounter::Unit::kBytes));
    }
    return std::string_view(static_cast<const char*>(buffer), fill);
  }

  uint64_t preadv(
      uint64_t offset,
      const std::vector<melon::Range<char*>>& buffers,
      filesystems::File::IoStats* stats = nullptr) const override {
    auto res = ReadFile::preadv(offset, buffers, stats);
    if (stats) {
      stats->addCounter(
          "read",
          RuntimeCounter(
              static_cast<int64_t>(res), RuntimeCounter::Unit::kBytes));
    }
    ++numIos_;
    return res;
  }

  // Asserts that 'bytes' is as would be read from 'offset'.
  void checkData(const void* bytes, uint64_t offset, int32_t size) {
    for (auto i = 0; i < size; ++i) {
      char expected = seed_ + offset + i;
      ASSERT_EQ(expected, reinterpret_cast<const char*>(bytes)[i])
          << " at " << offset + i;
    }
  }

  uint64_t memoryUsage() const override {
    POLLUX_NYI();
  }

  bool shouldCoalesce() const override {
    POLLUX_NYI();
  }

  int64_t numIos() const {
    return numIos_;
  }

  std::string getName() const override {
    return "<TestReadFile>";
  }

  uint64_t getNaturalReadSize() const override {
    POLLUX_NYI();
  }

 private:
  const uint64_t seed_;
  const uint64_t length_;
  std::shared_ptr<filesystems::File::IoStats> ioStats_;
  mutable std::atomic<int64_t> numIos_{0};
};

} // namespace kumo::pollux::dwio::common
